#include "vtkMeshExport.hpp"
#include "vtkMeshImport.hpp"
#include "vtkFieldExport_source.hpp"
#include "vtkFieldImport_source.hpp"
#include "sweSolver.hpp"

#include "meshTypeFunctions.hpp"

#include <catch2/catch_test_macros.hpp>

#include <cmath>

using namespace hms;
using namespace hms::vtk;
namespace xml = tinyxml2;

#define REQUIRE_RAINFALL(REQUIREMENT) \
readSourceFields(sourceDoc.get(), solver); \
std::visit( [&](auto& field){ \
	field.updateValues(); \
	REQUIRE((REQUIREMENT)); \
}, solver.r() )

TEST_CASE("Source field import"){
	auto test = []( auto&& meshDummy){
		using MeshType = remove_qualifiers<decltype(meshDummy)>;
		scalar dx {1.3}, dy {0.8};

		MeshVariant meshVar { uniMeshAs<MeshType>(3,4,dx,dy) };

		const MeshType& mesh { std::get<MeshType>(meshVar) };

		Time time {10};
		time.writeTimeInterval(1);

		SWESolver<MeshType> solver {mesh, time};

		solver.r() = SourceField<Uniform, MeshType>{mesh, Array1s{0}};
		solver.i() = SourceField<Uniform, MeshType>{mesh, Array1s{0}};

		solver.useRainfall(true);
		solver.useInfiltration(true);

		std::unique_ptr<xml::XMLDocument> sourceDoc {
			sourceFieldsToXML(solver)
		};

		// toStream( sourceDoc.get(), std::cout );

		xml::XMLElement* sourceFields {
			getChildNodeByName( sourceDoc.get(), sourceNode )
		};
		xml::XMLElement* rain { getChildNodeByAttributeValue(sourceFields,
			sourceChild, "Name", "Rainfall"
		) };

		SECTION("Update method: none"){
			SECTION("Uniform"){
				/* one component */
				rain->SetText(1);
				REQUIRE_RAINFALL(
					field.updateMethod() == UpdateMethod::none &&
					field.storage().isApprox(Array1s{1})
				);

				/* two components */
				rain->SetAttribute(nComponentsName, 2);
				rain->SetText("2 3");
				REQUIRE_RAINFALL(
					field.updateMethod() == UpdateMethod::none &&
					field.storage().isApprox(Array2s{2,3})
				);
			}
			SECTION("Non-uniform"){
				/* one component */
				rain->SetAttribute("isUniform", false);
				ArrayXXs vals0 (1,12);
				vals0 << 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11;
				matrixToText(rain, vals0);
				REQUIRE_RAINFALL(
					field.updateMethod() == UpdateMethod::none &&
					field.domainValues().isApprox(vals0)
				);

				/* two components */
				rain->SetAttribute(nComponentsName, 2);
				ArrayXXs vals1 (2,12);
				vals1.row(0) = vals0;
				vals1.row(1) = 12;
				matrixToText(rain, vals1);
				REQUIRE_RAINFALL(
					field.updateMethod() == UpdateMethod::none &&
					field.domainValues().isApprox(vals1)
				);
			}
		}
		SECTION("Update method: timeSeries"){
			rain->SetAttribute("updateMethod", "timeSeries");
			rain->SetAttribute("useInterpolation", false);
			TimeSeries timeSeries;
			SECTION("Uniform"){
				timeSeries = { {0, Array1s{0}}, {1, Array1s{1}} };
				writeTimeSeries(rain, timeSeries);
				REQUIRE_RAINFALL(
					!field.useInterpolation() &&
					field.updateMethod() == UpdateMethod::timeSeries &&
					field.domainValues().isApproxToConstant(0)
				);
			}
			rain->SetAttribute("useInterpolation", true);
			SECTION("Non-uniform"){
				rain->SetAttribute("isUniform", false);
				ArrayXXs vals (1,12);
				vals = 5;
				timeSeries = { {0, vals}, {1, vals+1} };
				writeTimeSeries(rain, timeSeries);
				REQUIRE_RAINFALL(
					field.useInterpolation() &&
					field.updateMethod() == UpdateMethod::timeSeries &&
					field.domainValues().isApprox(vals)
				);
			}
		}
		SECTION("Update method: function"){
			rain->SetAttribute("updateMethod", "function");
			SECTION("Uniform"){
				std::string functionStr {"1/2*g*t"};
				writeFunction(rain, functionStr, 0);
				REQUIRE_RAINFALL(
					field.updateMethod() == UpdateMethod::function &&
					field.functions()[0].functionString() == functionStr
				);
			}
			SECTION("Non-uniform"){
				rain->SetAttribute("isUniform", false);
				rain->SetAttribute(nComponentsName, 2);
				std::string functionStr0 {"x+y"}, functionStr1 {"z*t"};
				writeFunction(rain, functionStr0, 0);
				writeFunction(rain, functionStr1, 1);
				REQUIRE_RAINFALL(
					field.updateMethod() == UpdateMethod::function &&
					field.functions()[0].functionString() == functionStr0 &&
					field.functions()[1].functionString() == functionStr1
				);
			}
		}
	};

	SECTION("Mesh"){
		test( Mesh{} );
	}
	SECTION("UniMesh"){
		test( UniMesh{} );
	}
	SECTION("RectMesh"){
		test( RectMesh{} );
	}
	SECTION("StructMesh"){
		test( StructMesh{} );
	}
}

#undef REQUIRE_RAINFALL